# This file contains the code for the car control node.
import rclpy
from rclpy.node import Node
import numpy as np
from std_msgs.msg import Float32MultiArray
import serial
import time
from geometry_msgs.msg import Pose, Twist, Vector3
from nav_msgs.msg import Odometry
import tf2_ros
from tf2_ros import TransformBroadcaster
from geometry_msgs.msg import TransformStamped
import math


Speed_PID_P=0.6
Speed_PID_I=0.01

CAR_WHEEL_WIDTH=0.19    # 车轮左右间距
CAR_WHEEL_DISTANCE=0.13 # 车轮前后间距
APB=(CAR_WHEEL_DISTANCE / 2 + CAR_WHEEL_WIDTH / 2)

WHEEL_D=80
PULSE_PER_ROLL=780   # 编码器每转一圈的脉冲数
ODOMETER_EST_PULSE_PER_METER=PULSE_PER_ROLL/(math.pi*WHEEL_D)


class CarCtrl(Node):
    def __init__(self,ser:serial.Serial):
        super().__init__('car_ctrl')
        self.get_logger().info('Hi from car_ctrl node.')
        
        # 设置参数
        self.declare_parameter('wheel_radius', 0.04) # 轮子半径
        self.declare_parameter('car_wheel_width', 0.19) # 车轮左右间距
        self.declare_parameter('car_wheel_distance', 0.13) # 车轮前后间距
        self.declare_parameter('encoder_pulse_per_meter', 780) # 编码器每米脉冲数
        self.declare_parameter('odometer_est_pulse_per_meter', 780/(math.pi*80)) # 里程计每米脉冲数
        
###### 初始化变量 ######
    #cur为状态量，用于计算函数的读取
        self.cur_wheels_speed=[0.0, 0.0, 0.0, 0.0]  # 当前四轮线速度，单位为m/s
        self.cur_wheels_pwm=[0, 0, 0, 0]    # 当前四轮速度对应的PWM值
        self.cur_wheels_speed_target=[0.0, 0.0, 0.0, 0.0]   # 当前四轮目标速度，单位为 m/s
        # self.cur_wheels_speed_target_pwn=[0.0, 0.0, 0.0, 0.0]  # 四轮目标速度，单位PWM，暂时废弃
        self.cur_wheels_flag=[0, 0, 0, 0]  # 当前四轮正转反转标志位，0为停止，1为正转，-1为反转
        
        self.last_wheels_speed=[0.0, 0.0, 0.0, 0.0]  # 上一次四轮线速度，单位为m/s
        
        self.cur_car_odom=[]
        self.cur_car_speed=[0.0, 0.0, 0.0]  # 当前车速，单位为m/s, x:前+后-  y:左+右-  z:逆+顺-
        
        self.wspd_bias=[0.0, 0.0, 0.0, 0.0]
        self.wspd_lastbias=[0.0, 0.0, 0.0, 0.0]

        self.ser=ser
###### 初始化变量 ######

# 轮速发布者 #
        # 添加发布者 : wheels_speed_publisher
        self.wheels_speed_publisher_ = self.create_publisher(Float32MultiArray, 'wheels_speed', 10)
        self.wheels_speed_msg_ = Float32MultiArray()
        self.wheels_speed_msg_.data = [0.0, 0.0, 0.0, 0.0]
        # 添加定时器 : timer_period_ms
        self.timer_period_ms = 100
        self.timer = self.create_timer(self.timer_period_ms/1000.0, self.wheel_speed_publisher_timer_callback)
        
# odom发布者 #
        # 添加发布者 : odom_publisher
        self.odom_publisher = self.create_publisher(Odometry, 'odom', 10)
        # 创建定时器,每隔0.1秒发布一次里程计信息
        self.odom_timer = self.create_timer(0.1, self.calc_car_speed_and_publish_odom)
        self.odom_frame_id = 'odom'
        self.child_frame_id = 'base_link'
        self.odom_broadcaster = TransformBroadcaster(self)
        # 初始化里程计变量
        self.odom_x = 0.0
        self.odom_y = 0.0
        self.odom_theta = 0.0
        # 例子中使用的速度
        self.odom_vx = 0.1  # 例子中使用的线速度
        self.odom_vtheta = 0.1  # 例子中使用的角速度
        
# 轮速设置订阅者 #
        # 添加 wheels_speed_set :设置轮子速度(PWM值)
        self.wheels_speed_set_subscriber = self.create_subscription(Float32MultiArray, 'wheels_speed_set', self.set_wheels_speed_callback, 10)

# 整车速度设置订阅者 #
        # 添加 car_speed_set by (vx,vy,vz) :设置整车速度(m/s)
        # x:前+后-  y:左+右-  z:逆+顺-
        self.car_speed_set=[0,0,0]
        # self.car_speed_set_subscriber = self.create_subscription(Float32MultiArray, 'car_speed_set', self.set_car_speed_callback, 10)
        self.car_speed_set_subscriber = self.create_subscription(Twist,'cmd_vel',self.set_car_speed_callback,10)

# ========回调函数BEGIN========

    def calc_car_speed_and_publish_odom(self):
        '''
        计算车速并发布里程计信息
        '''
        for i in range(4):
            if self.cur_wheels_pwm[i]>1500:
                self.cur_wheels_flag[i]=1
            elif self.cur_wheels_pwm[i]<1500:
                self.cur_wheels_flag[i]=-1
            else:
                self.cur_wheels_flag[i]=0

            self.cur_wheels_speed[i]=self.cur_wheels_flag[i]*self.cur_wheels_speed[i]
        self.cur_car_speed=self.car_inverse_kinematics(self.cur_wheels_speed)
        self.odom_x+=self.cur_car_speed[0]*0.1
        self.odom_y+=self.cur_car_speed[1]*0.1
        self.odom_theta+=self.cur_car_speed[2]*0.1
        self.publish_odometry()


    def set_car_speed_callback(self, msg):
        '''
        设置整车速度 回调函数
        '''
        # 临时变量
        self.car_speed_set = msg.data
        self.get_logger().info('[Received] set_car_speed : %s' % self.car_speed_set)
        # self.set_car_speed(speed_value=self.car_speed_set)

    def set_wheels_speed_callback(self, msg):
        '''
        设置轮子速度 回调函数
        '''
        # 全局状态变量设置
        # self.cur_wheels_speed_target = msg.data
        # self.wheels_speed_set=self.cur_wheels_speed_target
        # 在此处更改cur_wheels_speed_pwn
        tmp_list=msg.data
        # self.cur_wheels_pwm= tmp_list
        self.get_logger().info('[Received] set_wheels_speed : %s' % tmp_list)
        self.set_motors_speed(speed_values=tmp_list)

    def wheel_speed_publisher_timer_callback(self):
        '''
        定时器回调函数，发布当前四轮速度信息
        '''
        # 全局状态变量更新
        self.last_wheels_speed = self.cur_wheels_speed
        self.cur_wheels_speed = self.get_speed_LAR(self.ser)
        # 更新轮速消息,设为当前速度
        self.wheels_speed_msg_.data = self.cur_wheels_speed
        # 发布 wheels_speed_msg_
        self.wheels_speed_publisher_.publish(self.wheels_speed_msg_)
        # self.get_logger().info('[Published] wheels_speed_msg_ : %s' % self.wheels_speed_msg_.data)
        
    def publish_odometry(self):
        '''
        发布里程计信息，计算在calc_car_speed_and_publish_odom函数中
        
        '''
        # Create a new message
        odom = Odometry()

        # Header
        odom.header.stamp = self.get_clock().now().to_msg()
        odom.header.frame_id = 'odom'
        odom.child_frame_id = 'base_link'

        # Position
        # Here you would normally get the position from a sensor or odometry system
        odom.pose.pose.position.x = self.odom_x  # Replace with actual x position
        odom.pose.pose.position.y = self.odom_y  # Replace with actual y position
        odom.pose.pose.position.z = 0.0  # Replace with actual z position

        # Orientation
        # Here you would normally get the orientation from a sensor or odometry system
        odom.pose.pose.orientation.x = 0.0  # Replace with actual x orientation
        odom.pose.pose.orientation.y = 0.0  # Replace with actual y orientation
        odom.pose.pose.orientation.z = 0.0  # Replace with actual z orientation
        odom.pose.pose.orientation.w = 1.0  # Replace with actual w orientation

        # Covariance (you can set these to any value)
        odom.pose.covariance = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
                                                0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
                                                0.0, 0.0, 0.0, 0.0, 0.0, 0.0]

        # Twist
        # Here you would normally get the velocity from a sensor or odometry system
        odom.twist.twist.linear.x = self.cur_car_speed[0]  # Replace with actual x linear velocity
        odom.twist.twist.linear.y = self.cur_car_speed[1]  # Replace with actual y linear velocity
        odom.twist.twist.linear.z = 0.0  # Replace with actual z linear velocity
        odom.twist.twist.angular.x = 0.0  # Replace with actual x angular velocity
        odom.twist.twist.angular.y = 0.0  # Replace with actual y angular velocity
        odom.twist.twist.angular.z = self.cur_car_speed[2]  # Replace with actual z angular velocity

        # Covariance (you can set these to any value)
        odom.twist.covariance = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
                                                0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
                                                0.0, 0.0, 0.0, 0.0, 0.0, 0.0]

        # Publish the message
        self.odom_publisher.publish(odom)

        # Publish the transform
        # t = TransformStamped()
        # t.header.stamp = odom.header.stamp
        # t.header.frame_id = odom.header.frame_id
        # t.child_frame_id = odom.child_frame_id
        # t.transform.translation.x = odom.pose.pose.position.x
        # t.transform.translation.y = odom.pose.pose.position.y
        # t.transform.translation.z = odom.pose.pose.position.z
        # t.transform.rotation = odom.pose.pose.orientation
        # self.tf_broadcaster.sendTransform(t)
        
# ========回调函数END========

    def set_car_speed(self, speed_value):
        '''
        设置整车速度
        Args:
            speed_value: 整车速度，单位为m/s，x:前+后-  y:左+右-  z:逆+顺-
        '''
        speed_mps=self.car_kinematics(speed_value[0],speed_value[1],speed_value[2]) # 计算四轮线速度 m/s
        
        self.set_motors_speed(speed_mps)
#{OK}
    def get_speed_PULSE(self,con:serial.Serial):
        '''
        获取100ms内轮子脉冲数
        Args:
            con: 串口对象

        Returns:
            speed: 四个轮子的脉冲数
        '''
        #脉冲数M
        M=0
        speed=[0,0,0,0]
        data=""
        #发送指令
        for i in range(6,10):
            if M>=0:
                cmdtxt=f"#{i:03d}PSPD!"
                con.write(cmdtxt.encode('utf-8'))
                #接收返回值
                data = con.readline().decode('utf-8')
                try:
                    M=int(data[data.find('=')+1:data.find('/')])
                    if i==7:
                        if self.cur_wheels_pwm[i-6]>1500:
                            M=(M+7)*0.5
                        elif self.cur_wheels_pwm[i-6]<1500:
                            M=(M-1)*0.5
                except ValueError:
                    M=0
            speed[i-6]=M     #Now set the speed value in pulse per second
        return speed
#{OK}
    def get_speed_RAD(self,con:serial.Serial):
        '''
        获取轮子角速度,单位为rad/s
        
        Args:
            con: 串口对象

        Returns:
            speed: 四个轮子的速度值，单位为rad/s
        '''
        speed=self.get_speed_PULSE(con)
        for i in range(4):
            if speed[i]>0:
                speed[i]=speed[i]*10/780    #Now set the speed value in rad/s
            #speed[i-6]=M     #Now set the speed value in pulse per second
        return speed
#{OK}
    def get_speed_LAR(self,con:serial.Serial):
        '''
        获取轮子线速度,单位为m/s
        
        Args:
            con: 串口对象

        Returns:
            speed: 四个轮子的速度值，单位为m/s
        '''
        #set pi
        PI=3.14159265358979
        #轮子直径80mm
        D=80
        #脉冲数M
        M=0
        speed=self.get_speed_PULSE(con)
        for i in range(4):
            if speed[i]>0:
                M=speed[i]
            speed[i]=PI*D*M*10/780 
            # convert to m/s
            speed[i]=round(speed[i]/1000,2)
        return speed
#{OK}
    def num_mapping(self,num, a_min, a_max, b_min, b_max):
        """
        将整数从a_min到a_max线性映射到b_min到b_max
        Args:
            num: 整数
            a_min: 原区间最小值
            a_max: 原区间最大值
            b_min: 目标区间最小值
            b_max: 目标区间最大值

        Returns:
            映射后的整数值
        """
        a_range = a_max - a_min
        b_range = b_max - b_min
        scaled_num = float(num - a_min) / float(a_range)
        return int(b_min + (scaled_num * b_range))
#{OK}
    def cmd_set_motor_speed(self,motor_id, speed,con_time=0):
        '''
        获取指令字符串，设置单个轮子的速度(通过PWM)
        
        Args:
            motor_id: 轮子的ID, 范围为 006,007,008,009
            speed: 轮子的PWM值，范围为 500-2500
            con_time: 延时时间，范围为 0-9999

        Returns:
            cmd_txt: 指令字符串
        '''

        cmd_txt="#{:03d}P{:04d}T{:04d}!".format(motor_id, speed,con_time)
        return str(cmd_txt)
#{OK}
    def send_command(self,con, cmd_txt):
        '''
        发送指令到串口
        Args:
            con: 串口对象
            cmd_txt: 指令字符串（未编码）
        Returns:
            None
        
        '''
        con.write(cmd_txt.encode('utf-8'))
        # receive return value
        # data = con.readline().decode('utf-8')
        # print(data)
#{OK}
    def set_motors_speed(self,motor_id=6,speed_values=[0,0,0,0]):
        '''
        设置四个轮子的速度(通过PWM)，变量cur_wheels_pwm在此更新
        
        Args:
            motor_id: 轮子的ID, 范围为 006,007,008,009
            speed_values: 四个轮子的PWM值，范围为 500-2500

        Returns:
            None
        '''
        cmd_txt="{"
        for i in range(4):
            self.cur_wheels_pwm[i]=speed_values[i]
            cmd_txt+=self.cmd_set_motor_speed(motor_id+i,round(speed_values[i]))
        # send command to serial port
        cmd_txt+="}"
        self.send_command(self.ser, cmd_txt)
        # test print
        # print(cmd_txt)
#{OK}
    def car_kinematics(self,vx,vy,vth):
        '''
        小车运动学正解，通过底盘速度计算四个轮子的速度
        Args:
            vx: 前后速度
            vy: 左右速度
            vth: 转角速度

        Returns:
            speed: 四个轮子的速度值，单位为m/s
        '''
        speed=[0.0,0.0,0.0,0.0]
        speed[1]=vx+vy-vth*APB
        speed[0]=-vx+vy+vth*APB
        speed[2]=-vx+vy-vth*APB
        speed[3]=vx+vy+vth*APB
        for i in range(4):
            speed[i]=round(speed[i],2)  # 速度保留两位小数
        return speed

# 整车移动量转换为单轮Linear速度  x:前+后-  y:左+右-  z:逆+顺-
    # def move_transform(vx, vy, vz):
    #     '''
    #     将整车的移动量转换为四个轮子的目标速度,单位为m/s

    #     Args:
    #         vx: 前后速度
    #         vy: 左右速度
    #         vz: 方向速度

    #     Returns:
    #         target_speed: 四个轮子的目标速度
    #     '''
    #     APB=(CAR_WHEEL_DISTANCE / 2 + CAR_WHEEL_WIDTH / 2)
    #     # 计算四个轮子的目标速度
    #     a = vx - vy - vz * APB
    #     b = vx + vy + vz * APB
    #     c = vx + vy - vz * APB
    #     d = vx - vy + vz * APB
    #     # 计算每个轮子的目标速度
    #     target_speed = [a, -1*b, c, -1*d]
    #     return target_speed
#{OK}
    def car_inverse_kinematics(self,speed):
        vx = (speed[1] - speed[0] - speed[2] + speed[3]) / 4
        vy = (speed[1] + speed[0]) / 2
        vth = -(speed[1] - speed[0] + speed[2] - speed[3]) / (4 * APB)
        return vx, vy, vth


    def set_speed_PID(self,target_speed):

        # PID control
        for i in range(4):
            if target_speed[i]==0:
                self.wheels_speed_set[i] = 1500
            else:
                # calculate error
                error = target_speed[i] - self.cur_wheels_speed[i]
                # calculate P term
                P_term = Speed_PID_P * (error-self.wspd_lastbias[i])
                # calculate I term
                I_term = Speed_PID_I * self.cur_wheels_speed[i]
                # calculate D term
                D_term = 0
                # calculate speed
                speed = P_term + I_term + D_term
                # set speed
                self.wheels_speed_set[i] += speed*10
                # limit speed
                if self.wheels_speed_set[i]>2500:
                    self.wheels_speed_set[i]=2500
                elif self.wheels_speed_set[i]<500:
                    self.wheels_speed_set[i]=500
                # update last bias
                self.wspd_lastbias[i] = error
        # send command to serial port
        self.set_motors_speed(speed_values=self.wheels_speed_set)

    def fun1():
        pass


def main():
    time.sleep(1)
    # 打开串口
    ser=None
    try:
        ser = serial.Serial('/dev/ttyUSB0', 115200, timeout=0.1)
    except Exception as e:
        print(f"Error opening serial port: {e},Retrying...")
        time.sleep(2)
    while ser is None or not ser.is_open:
        try:
            ser = serial.Serial('/dev/ttyUSB0', 115200, timeout=0.1)
        except Exception as e:
            print(f"Error opening serial port: {e},Retrying...")
            time.sleep(2)
    s=ser.readline()
    while s.decode().find("[FinishInit]")==-1:
        s=ser.readline()

    # ROS初始化
    rclpy.init()
    car_ctrl = CarCtrl(ser)
    rclpy.spin(car_ctrl)
    car_ctrl.destroy_node()
    rclpy.shutdown()


if __name__ == '__main__':
    main()
